#include <stdlib.h>
#include <time.h>
#include <assert.h>
#include <sys/time.h>
#include <vector>
#include <omp.h>
#include "define.h"
#include "expVec.cc"

#define DIM 3

#define DEFAULT_NUMBER_OF_POINTS 256
#define DEFAULT_NUMBER_OF_BOXES 1024

#define X(s,k) (s)[(k)*DIM]
#define Y(s,k) (s)[(k)*DIM+1]
#define Z(s,k) (s)[(k)*DIM+2]
#define INVSQRT(R) {if( (R)!=0 ) (R)=1.0/(R);}

#ifdef USE_PD
#define _real double
#define quad __m128d
#else
#define _real float
#define quad __m128
#endif

#define BLK 64
using std::min;
using namespace std;
typedef complex<_real> dcomp;

#ifdef USE_COMP_COEF
#define k_coeff dcomp
#else
#define k_coeff _real
#endif

double timediff(struct timeval start, struct timeval stop){
	return (stop.tv_sec -  start.tv_sec) + 1.0e-6 * (stop.tv_usec  - start.tv_usec);
}

/* ************************************************** */
void compexp_Direct (const int ns, const int nt,  _real const src[], _real const trg[], dcomp const den[], dcomp pot[],  k_coeff kernel_coef )
{
	dcomp mult = dcomp(0.0, -1.0);
	dcomp mult_coeff = kernel_coef * mult;
	for (int t=0; t<nt; t++){
		dcomp p= dcomp(0,0);
		_real tx=X(trg,t);
		_real ty=Y(trg,t);
		_real tz=Z(trg,t);
		for (int s=0; s<ns; s++){
			_real rx=X(src,s)-tx;
			_real ry=Y(src,s)-ty;
			_real rz=Z(src,s)-tz;
			_real rho = rx*rx + ry*ry + rz*rz;
			rho = sqrt(rho);
			dcomp d = den[s];
			dcomp Kernel = exp(mult_coeff * rho);
			INVSQRT(rho);
			p += (Kernel * d) * rho;
		}
		pot[t] = p;
	}
	return;
}

/* ************************************************** */
void compexp_blocked (const int ns, const int nt,  _real const src[], _real const trg[], dcomp const den[], dcomp pot[], const k_coeff kernel_coef )
{
	dcomp mult = dcomp(0., -1.);
	dcomp mult_coeff = kernel_coef * mult;
	for (int tt=0; tt<nt; tt+=BLK) {
		for (int ss=0; ss<ns; ss+=BLK) {
			for (int t=tt; t<min(nt, tt+BLK); t++){
				_real tx=X(trg,t);
				_real ty=Y(trg,t);
				_real tz=Z(trg,t);
				int s;
				dcomp p= dcomp(0,0);
				for (s=ss; s<min(ns, ss+BLK); s++){		
					_real rx=X(src,s)-tx;
					_real ry=Y(src,s)-ty;
					_real rz=Z(src,s)-tz;
					_real rho = rx*rx + ry*ry + rz*rz;
					rho = sqrt(rho);
					dcomp d = den[s];
					dcomp Kernel = exp(mult_coeff*rho);
					INVSQRT(rho);
					p += (Kernel * d)* rho;
				}
				pot[t] += p;
			}
		}
	}
	return;
}

/* ************************************************** */
void compexp_Shuffle_SSE (const int ns, const int nt,  _real const src[], _real const trg[], dcomp const den[], dcomp pot[], const k_coeff kernel_coef, _real workspace[]  ) {
	
	//1. reshuffle memory
	for (int k =0;k<ns;k++){
		workspace[k] = X(src,k);
		workspace[ns+k] = Y(src,k);
		workspace[ns*2+k] = Z(src,k);
	}
	int offset = ns*DIM;
	_real* target_workspace = workspace + offset;
	for (int k=0;k<nt;k++){
		target_workspace[k] = X(trg,k);
		target_workspace[nt+k] = Y(trg,k);
		target_workspace[nt*2+k] = Z(trg,k);
	}
	
	//2. perform caclulation 
	compexp_SSE<_real, quad, k_coeff>(ns,nt,workspace,workspace+ns,workspace+2*ns,workspace+3*ns,workspace+3*ns+nt,workspace+3*ns+2*nt,den,pot,kernel_coef);
	return;
}

/* ************************************************** */
void compexp_after_Shuffle (const int ns, const int nt,  _real const xs[], _real const ys[], _real const zs[], _real const xt[], _real const yt[], _real const zt[],
	dcomp const den[], dcomp pot[], const k_coeff kernel_coef ) {
	dcomp mult = dcomp(0.0, -1.0);
	dcomp mult_coeff = kernel_coef * mult;
	for(int t=0; t < nt; t++) {
		dcomp p = dcomp(0,0);
		_real tx = xt[t];
		_real ty = yt[t];
		_real tz = zt[t];
		for(int s=0; s < ns; s++) {
			_real rx = xs[s] - tx;
			_real ry = ys[s] - ty;
			_real rz = zs[s] - tz;
			_real rho = rx*rx + ry*ry + rz*rz;
			rho = sqrt(rho);
			dcomp d = den[s];
			dcomp Kernel = exp(mult_coeff * rho);
			INVSQRT(rho);
			p += (Kernel * d) * rho;
		}
		pot[t] = p;
	}
	return;
}
/* ************************************************** */
void compexp_Shuffle_Direct (const int ns, const int nt,  _real const src[], _real const trg[], dcomp const den[], dcomp pot[], const k_coeff kernel_coef, _real workspace[]  ) {
	
	//1. reshuffle memory
	for (int k =0;k<ns;k++){
		workspace[k] = X(src,k);
		workspace[ns+k] = Y(src,k);
		workspace[ns*2+k] = Z(src,k);
	}
	int offset = ns*DIM;
	_real* target_workspace = workspace + offset;
	for (int k=0;k<nt;k++){
		target_workspace[k] = X(trg,k);
		target_workspace[nt+k] = Y(trg,k);
		target_workspace[nt*2+k] = Z(trg,k);
	}
	
	//2. perform caclulation 
	compexp_after_Shuffle(ns, nt,workspace,workspace+ns,workspace+2*ns,workspace+3*ns,workspace+3*ns+nt,workspace+3*ns+2*nt,den,pot,kernel_coef);
	return;
}

/* ************************************************** */
int main(int argc, char **argv){

	struct timeval start, stop;

	int n = DEFAULT_NUMBER_OF_POINTS;

	printf("Usage:\n newFMM [number_of_points_per_box] [number_of_boxes]\n\n");
	if (argc>1) n = atoi(argv[1]);
	int nt=DEFAULT_NUMBER_OF_BOXES;
	if (argc>2) nt = atoi(argv[2]);
#ifdef USE_COMP_COEF
	k_coeff kercof= dcomp(10.0, -10.0);
#else
	k_coeff kercof = 10.0;
#endif
	//1. allocate memory for sources and targets.
	_real *src = new _real[n*DIM*nt];
	_real *trg = new _real[n*DIM*nt];
	dcomp *den = new dcomp[n*nt];
	
	//2. initialize particle positions and density
	_real h =1.0/n/nt;
	for (int k=0;k<n*nt;k++){
		_real st=h*k;
		X(src,k) = st;Y(src,k) = st;Z(src,k) = st;
		X(trg,k) = 2+st;  Y(trg,k) = 2+st; Z(trg,k) = 2+st;
		den[k] = dcomp(h,h);
	}

	//3. call direct interactions
	
	double t, t1;
	//3.1 simple matrix-fee pairwise interactions
	dcomp *pot1 = new dcomp[n*nt];
	gettimeofday(&start, 0);
	for( int k=0; k<nt; k++)
	compexp_Direct(n,n,src+k*n*DIM,trg+k*n*DIM,den+k*n,pot1+k*n,kercof);
	gettimeofday(&stop, 0);
	t = timediff(start, stop);
	printf("Time for sequential %f\n\n",t);	
	int k;
	//3.2 simple matrix-fee pairwise interactions with OpenMP
	dcomp *pot2 = new dcomp[n*nt];
	gettimeofday(&start, 0);
	
	#pragma omp parallel for private(k) shared(n,src,trg,den,pot2, nt) schedule(dynamic)
	for(k=0; k<nt; k++) {
		compexp_Direct(n,n,src+k*n*DIM,trg+k*n*DIM,den+k*n,pot2+k*n,kercof);	
	}
	gettimeofday(&stop, 0);
	t1 = timediff(start, stop);
	
	printf("Time for OpenMp %f\n",t1);	
	printf("Speedup with OMP = %f\n\n", t/t1);
	
	//3.3 Shuffle and direct computation
	dcomp *pot3 = new dcomp[n*nt];
	gettimeofday(&start, 0);
	int threads_ = omp_get_max_threads();
	_real *workspace_ = new _real[2*DIM*n*threads_];
	#pragma omp parallel for private(k) shared(n,src,trg,den,pot3, nt) schedule(dynamic)
	for(k=0; k<nt; k++) {
		int thread = omp_get_thread_num();
		compexp_Shuffle_Direct(n, n, src+k*n*DIM, trg+k*n*DIM, den+k*n, pot3+k*n, kercof, workspace_+2*DIM*n*thread);
	}
	gettimeofday(&stop, 0);
	t1 = timediff(start, stop);
	
	printf("Time for OMP + shuffle %f\n",t1);	
	printf("Speedup with OMP + shuffle = %f\n\n", t/t1);

	//3.4 simple matrix-fee pairwise interactions with OpenMP and SSE
	dcomp *pot4 = new dcomp[n*nt];	
	gettimeofday(&start, 0);
	init_SSE<_real>();
	int threads = omp_get_max_threads();
	_real *workspace = new _real[2*DIM*n*threads];
	#pragma omp parallel for private(k) shared(n,src,trg,den,pot4, nt) schedule(dynamic)
	for(k=0; k<nt; k++) {
		int thread = omp_get_thread_num();
		compexp_Shuffle_SSE(n, n, src+k*n*DIM, trg+k*n*DIM, den+k*n, pot4+k*n, kercof, workspace+2*DIM*n*thread);
	}
	gettimeofday(&stop, 0);
	t1 = timediff(start, stop);
	
	printf("Time for OMP+SSE %f\n",t1);	
	printf("Speedup with OMP + SSE = %f\n\n", t/t1);
	

	
	for(int i =0; i < nt*n; i++) {
			if(fabs(real(pot1[i] - pot4[i])) > 0.0000001 || fabs(imag(pot1[i] - pot4[i])) > 0.0000001) 
				cout << i << " " << pot1[i] << "  " <<  pot4[i] << endl;	
	}
	
	return 0;
	
}

